#!/usr/bin/env python
# -*- encoding: utf-8 -*-


import elasticsearch.helpers
from elasticsearch import Elasticsearch
from pyspider.database.base.groupdb import GroupDB as BaseGroupDB


class GroupDB(BaseGroupDB):
    __type__ = 'group'

    def __init__(self, hosts, index='pyspider'):
        self.index = index
        self.es = Elasticsearch(hosts=hosts)

        self.es.indices.create(index=self.index, ignore=400)
        if not self.es.indices.get_mapping(index=self.index, doc_type=self.__type__):
            self.es.indices.put_mapping(index=self.index, doc_type=self.__type__, body={
                "_all": {"enabled": True},
                "properties": {
                    "group_keyword": {"type": "string", "index": "analyzed"},
                    "resultid": {"type": "string", "index": "analyzed"},
                    "project": {"type": "string", "index": "not_analyzed"},
                }
            })

    def insert_group(self, group_keyword, resultid, title, url, project, date):
        obj = {
            'group_keyword': group_keyword,
            'resultid': resultid,
            'title': title,
            'url': url,
            'project': project,
            'date': date,
        }
        return self.es.index(index=self.index, doc_type=self.__type__,
                        body=obj, id='%s:%s' % (group_keyword, resultid))

    def refresh(self):
        """
        Explicitly refresh one or more index, making all operations
        performed since the last refresh available for search.
        """
        self.es.indices.refresh(index=self.index)

    def drop(self, project):
        self.refresh()
        for record in elasticsearch.helpers.scan(self.es, index=self.index, doc_type=self.__type__,
                                                 query={'query': {'term': {'project': project}}},
                                                 _source=False):
            self.es.delete(index=self.index, doc_type=self.__type__, id=record['_id'])
